/**
 * User-level library implementation of a memory sandbox
 *
 * Single Author info:  trhardy  Trevor R Hardy
 *
 * Group info:
 * trhardy  Trevor R Hardy
 * payyapp  Prathibha Ayyappan
 * vsubram3  Vidhatri Subramanya
 */

#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <ucontext.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <unistd.h>
#include <errno.h>

void sandbox_protection_handle(int sig, siginfo_t *signalInfo, struct ucontext *contextInfo);

unsigned int (*addr_ReadString)(char *s, int r);

unsigned int (*addr_main)(int argc, char *argv[]);

extern void readString(char *s, int r); 

extern void main(int argc, char *argv[]);

int ret_addr;

/**
 *  Initializes the memory sandbox to detect buffer overflow attacks by setting up signal handling for SIGSEGV and memory protection on the readString and main pages
 *
 */
void init_sandbox(){
	printf("Entering init_sandbox!\n");

	/* Set the handler for SIGSEGV to be sandbox_protection_handle */
	struct sigaction new_handler;
	new_handler.sa_sigaction = (void*)sandbox_protection_handle;
	sigemptyset(&new_handler.sa_mask);
	sigaddset(&new_handler.sa_mask, SIGSEGV);
	new_handler.sa_flags = 0;
	new_handler.sa_flags = new_handler.sa_flags | SA_SIGINFO | SA_RESTART;
	if(sigaction(SIGSEGV, &new_handler, NULL) == -1){
		printf("Sigaction FAILED!\n");
	}
	
	/* Unmask SIGSEGV */
	sigset_t new_mask;
	sigemptyset(&new_mask);
	sigaddset(&new_mask, SIGSEGV);
	if(sigprocmask(SIG_UNBLOCK, &new_mask, NULL) != 0){
		printf("Sigprocmask FAILED!\n");
	}	
	
	/* Set only read permissions on the readString function page(s) */
	addr_ReadString = (void*) readString;
	if(mprotect((void*)((unsigned long)addr_ReadString & 0xfffffffffffff000), 4, PROT_READ) == -1){
		perror("mprotect failed!");
	}
	
	/* Set read and execute permissions on the main function page(s) */
	addr_main = (void*) main;
	if(mprotect((void*)((unsigned long)addr_main & 0xfffffffffffff000), 4, PROT_READ | PROT_EXEC) == -1){
		perror("mprotect failed!");
	}
}

/**
 *  Handler for SIGSEGV that detects attempts at memory manipulation using buffer overflow attacks
 *
 *@param sig  This indicates the signal that is being caught
 *@param signalInfo  This struct contains information about the signal being handled, most importantly where the fault occurred
 *@param contextInfo  This struct contains information oabout the system state at the time of the fault, most importantly the stack pointer at the time of the fault
 */
void sandbox_protection_handle(int sig, siginfo_t * signalInfo, struct ucontext *contextInfo){
	printf(" entering handle\n");

	/* Get address where seg fault occurred */
	void *segAddr = signalInfo->si_addr;
	
	/* Get register state from when seg fault occurred, so that SP can be located */
  //  ucontext_t *regInfo = (ucontext_t *) contextInfo;	
	/*mcontext_t mContext = regInfo->uc_mcontext;
           gregset_t *registers = mContext.gregs;*/     //--- doubt

		   
	/* If the signal came upon entering readString, then protect main and unprotect readString */
	if(((long)segAddr & 0xfffffffffffff000) == ((long)addr_ReadString & 0xfffffffffffff000)){
		ret_addr = *(int *)(contextInfo-> uc_mcontext.gregs[15]); 	// -- doubt
		mprotect((void*)((long)addr_main & 0xfffffffffffff000), 4, PROT_READ);
		mprotect((void*)((long)addr_ReadString & 0xfffffffffffff000), 4, PROT_EXEC | PROT_READ);
	} else{
	    /* If the signal came while entering main, then protect readString, and unprotect main */
		if(((long)segAddr & 0xfffffffffffff000) == ((long)addr_main & 0xfffffffffffff000)){
			/* If the return address originally recorded upon entrance to the readString function has changed, then overflow occurred and print warning message and exit */
			if(ret_addr != (long)segAddr) {
				printf("malicious buffer overflow detected, will exit!\n");
				exit(0);
			}

			mprotect((void*)((long)addr_main & 0xfffffffffffff000), 4, PROT_EXEC | PROT_READ);
			mprotect((void*)((long)addr_ReadString & 0xfffffffffffff000), 4, PROT_READ);
		}
	}
}
